import argparse
from ipaddress import IPv4Address

from process_log_file import rank_origins, query_address_asn
from analyze_ssfnet_output import get_forwarding_table
from TrafficGenerator.traceroute_utils import tracert_path, unix_process_tracert_path, is_localhost_address

def my_tracert_path(hostname, unix_process_cmdline):
    if unix_process_cmdline:
        return unix_process_tracert_path(hostname, unix_process_cmdline)
    return tracert_path(hostname)

def ratio(file_type, input_file, cut, ssfnet_output_file, source_AS, output_file, unix_process_cmdline, *args, **kwargs):
    forwarding_table = get_forwarding_table(ssfnet_output_file, source_AS)

    ratio_table = {}
    AS_ranking = rank_origins(file_type, input_file, None, cut, False)


    for asn_record in AS_ranking:
        # [1] is for the value index in the (key, value) tuple.
        # [2] is for the IP list in the value tuple which looks like this (name, # of packets, IP list)
        # [0] is for the FIRST IP in the IP list.
        first_ip = asn_record[1][2][0]
        if is_localhost_address(IPv4Address(first_ip)):
            continue
        real_ip_path = my_tracert_path(first_ip, unix_process_cmdline)

        # Something went wrong. Try a different IP
        if len(real_ip_path) >= 30 and len(asn_record[1][2]) > 1:
            if is_localhost_address(IPv4Address(asn_record[1][2][1])):
                continue
            real_ip_path = my_tracert_path(asn_record[1][2][1], unix_process_cmdline)

        # Translate to ASNs
        real_AS_path = []
        for number, ip in real_ip_path:
            if is_localhost_address(IPv4Address(ip)):
                continue
            try:
                real_AS_path.append((number, query_address_asn(ip)))
            except Exception:
                print('No ASN found for host {0}'.format(ip))

        asn = asn_record[0]
        if asn in forwarding_table:
            # Appending the source as to the forwarding table, so itll match TRACERT results
            ratio_table[asn]=(real_AS_path, [source_AS] + forwarding_table[asn])

        hops = {}
        # Collect the hops in each AS
        for _asn, paths in ratio_table.items():
            real_AS_path_copy = paths[0].copy()
            forwarding_list = paths[1]
            while real_AS_path_copy:
                # 0 is the first item, 1 is the AS entry, 0 is the asn
                first_asn = int(real_AS_path_copy[0][1][0])
                hops_in_this_as = len(list(filter(lambda x: x[1][0] == first_asn, real_AS_path_copy)))
                if first_asn not in hops:
                    hops[first_asn] = [hops_in_this_as]
                else:
                    hops[first_asn].append(hops_in_this_as)
                # Remove this asn from the traceroute list
                real_AS_path_copy = real_AS_path_copy[hops_in_this_as:]


        with open(output_file, 'w') as outfile:
            for asn, paths in ratio_table.items():
                real_AS_path = paths[0]
                forwarding_list = paths[1]
                outfile.write('ASN: {0}\r\n'.format(asn))
                aggregated_list = []
                seen = []
                for number, AS in real_AS_path[:-1]:
                    if not AS[0] in seen:
                        aggregated_list.append((number, AS))
                        seen.append(AS[0])
                # Always write the last one
                aggregated_list.append(real_AS_path[-1])

                outfile.write('Aggregated AS Path by tracert: {0}\r\n'.format(aggregated_list))
                outfile.write('AS Path by simulation: {0}\r\n'.format(forwarding_list))
                outfile.write('Ratio: {0}'.format(len(real_AS_path) / len(forwarding_list)))
                outfile.write('\r\n\r\n')
            outfile.write('Ratios per AS:\r\n\r\n')
            for asn, hop_list in hops.items():
                outfile.write('Hops collected for asn {0}: {1}\r\n'.format(asn, hop_list))
                outfile.write('Average hop penalty: {0}\r\n'.format(sum(hop_list)/len(hop_list)))



if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    subparsers = parser.add_subparsers()

    # Simple ratio calculation
    incoming_parser = subparsers.add_parser('ratio', help='Calculate a simple ratio')
    incoming_parser.set_defaults(func=ratio)

    parser.add_argument("file_type", choices=['tcpdump', 'firewall'])
    parser.add_argument("input_file",)
    parser.add_argument('--cut', default=30, help='Show only this amount of ranked ASes')
    parser.add_argument("ssfnet_output_file", help="A console output file of SSFNET")
    parser.add_argument("source_AS", help="The source AS asn")
    parser.add_argument("output_file", help="Results file")
    parser.add_argument("-u", '--unix_process_cmdline', help="Use this command line to open a unix shell process for tracerouting")

    args = parser.parse_args()
    args.func(**vars(args))